Ground state simulation and quench of the Bose-Hubbard model in 2D.

# pylint: disable=invalid-name

import qtealeaves as qtl
from qtealeaves.models import get_bose_hubbard_2d


def main(tn_type=5, input_folder=None, output_folder=None, timesteps=160):
    """
    Main method for the dynamics of the Bose-Hubbard model in 2d.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.

    timesteps : int, optional
        NUmber of timesteps. Default to 160.
    """
    if input_folder is None:
        input_folder = lambda params: "BOHU/input_L%d" % (params["L"])
    if output_folder is None:
        output_folder = lambda params: "BOHU/output_L%d" % (params["L"])

    model, my_ops = get_bose_hubbard_2d()

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(
        max_iter=7, max_bond_dimension=16
    )
    my_obs = qtl.observables.TNObservables()

    # Dynamics
    quench = qtl.DynamicsQuench([0.05] * timesteps, time_evolution_mode=2)
    quench["J"] = lambda tt, params: params["J"] * (1 - tt / 8.0)

    simulation = qtl.QuantumGreenTeaSimulation(
        model,
        my_ops,
        my_conv,
        my_obs,
        tn_type=tn_type,
        folder_name_input=input_folder,
        folder_name_output=output_folder,
        has_log_file=False,
        store_checkpoints=False,
    )

    params = []

    params.append(
        {
            "L": 4,
            "U": 1.0,
            "J": 4.0,
            "Quenches": [quench],
            "exclude_from_hash": ["Quenches"],
        }
    )

    simulation.run(params, delete_existing_folder=True)

    for elem in params:
        e0_tn = simulation.get_static_obs(elem)["energy"]

        print("Energies ground state TTN", e0_tn)

    print(f"\nExample `{__file__}` ran successfully; no asserts implemented.")

    return


if __name__ == "__main__":
    main()